import numpy as np
import matplotlib.pyplot as plt

# Parameters
n_values = [64, 128]  # Grid sizes
omega = 2 / 3  # Weight for weighted Jacobi
nu_pairs = [(0, 0), (0, 2), (1, 1), (2, 0), (2, 2), (4, 0)]  # (ν1, ν2) pairs

# Function to compute damping coefficients
def damping_coefficients(n, omega, nu1, nu2):
    k_values = np.arange(1, n)  # Wavenumbers from 1 to n-1
    lambda_k = 1 - 2 * omega * np.sin(k_values * np.pi / (2*n)) ** 2  # Eigenvalues of weighted Jacobi
    s_k = np.sin(k_values * np.pi / (2*n)) ** 2 # Smoothing coefficients
    c_k = np.cos(k_values * np.pi / (2*n)) ** 2  # Coarse-grid coefficients

    # Damping coefficients
    c1 = s_k * lambda_k ** (nu1 + nu2)  # c1 = λ_k^(ν1 + ν2) s_k
    c2 = s_k * lambda_k ** nu1 * lambda_k[::-1] ** nu2  # c2 = s_k λ_k^ν1 λ_k'^ν2
    c3 = c_k * lambda_k[::-1] ** nu1 * lambda_k ** nu2  # c3 = c_k λ_k'^ν1 λ_k^ν2
    c4 = c_k * lambda_k[::-1] ** (nu1 + nu2)  # c4 = c_k λ_k'^(ν1 + ν2)

    return k_values, c1, c2, c3, c4

# Plot damping coefficients for n = 64 and n = 128 separately
for n in n_values:
    plt.figure(figsize=(18, 6))  # Create a new figure for each n

    for j, (nu1, nu2) in enumerate(nu_pairs):
        k_values, c1, c2, c3, c4 = damping_coefficients(n, omega, nu1, nu2)

        plt.subplot(2, 3, j + 1)  # Arrange subplots in 2 rows and 3 columns
        plt.plot(k_values, c1, 'b-', label='$c_1$')
        plt.plot(k_values, c2, 'r--', label='$c_2$')
        plt.plot(k_values, c3, 'm--', label='$c_3$')
        plt.plot(k_values, c4, 'g-', label='$c_4$')

        plt.title(f'Damping Coefficients for $n = {n}$, $(\\nu_1, \\nu_2) = ({nu1}, {nu2})$')
        plt.xlabel('Wavenumber $k$')
        plt.ylabel('Damping Coefficient')
        plt.legend()
        plt.grid(True)

    plt.tight_layout()
    plt.savefig(f'9.41_{n}.png', dpi=300, bbox_inches='tight')
    plt.show()